import os

from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate, \
    MessagesPlaceholder
from langchain.schema import ChatMessage, HumanMessage, AIMessage

prompt = "May the {subject} be with you"
chat_message_prompt = ChatMessagePromptTemplate.from_template(role="Jedi", template=prompt)
chat_message_prompt.format(subject="force")
ChatMessage(content="May the force be with you", additional_kwargs={}, role="Jedi")

human_prompt = "Summarize our conversation so far in{word_count} words."
human_message_template = HumanMessagePromptTemplate.from_template(human_prompt)
chat_prompt = ChatPromptTemplate.from_messages(
    [MessagesPlaceholder(variable_name="conversation"), human_message_template])
human_message = HumanMessage(content="What is the best way to learn programming?")
ai_message = AIMessage(content="""\
1. Choose a programming language: Decide on a programming language that you want to learn.

2. Start with the basics: Familiarize yourself with the basic programming concepts such as variables, data types and control structures.

3. Practice, practice, practice: The best way to learn programming is through hands-on experience\
""")
chat_prompt.format_prompt(conversation=[human_prompt, ai_message], word_count="10").to_messages()
messages = [HumanMessage(content='What is the best way to learn programming?', additional_kwargs={}),
            AIMessage(
                content='1. Choose a programming language: Decide on a programming language that you want to learn. '
                        '\n\n2. Start with the basics: Familiarize yourself with the basic programming concepts such '
                        'as variables, data types and control structures.\n\n3. Practice, practice, practice: The best '
                        'way to learn programming is through hands-on experience', additional_kwargs={}),
            HumanMessage(content='Summarize our conversation so far in 10 words.', additional_kwargs={})]
chat = ChatOpenAI(temperature=0)
chain = LLMChain(llm=chat, prompt=chat_prompt)
response = chain.run(messages)
